/*
 * Slicer.cpp
 *
 *  Created on: Nov 16, 2010
 *      Author: xavier
 */

#include "xiaowei/Slicer.h"

#include "llvm/Support/raw_ostream.h"
#include "llvm/Instructions.h"
#include <iostream>
#include <queue>

using namespace llvm;
using namespace klee;
using namespace std;

long Slicer::addKinsToDB(KInstruction * kins) {
	std::map<KInstruction *, long>::iterator it = kins_db.find(kins);
	// if not found
	long instance = 0;
	if (it == kins_db.end()) {
		kins_db.insert(pair<KInstruction *, long> (kins, 0));
	} else {
		instance = ++kins_db[kins];
	}
	return instance;
}

void Slicer::addToTrace(TraceEntry * entry) {
	trace.push_back(entry);
}

void Slicer::printInsDB() {
	std::map<KInstruction *, long>::iterator it;
	for (it = kins_db.begin(); it != kins_db.end(); it++) {
		llvm::errs() << it->first->inst << ": " << *(it->first->inst)
				<< " instance: " << it->second << "\n";
	}
}

int Slicer::mayWriteF(TraceEntry * te) {
	return true;
}

int Slicer::isConcrete(TraceEntry * te) {
	if (te->kins->inst->getOpcode() == Instruction::Store
			|| te->kins->inst->getOpcode() == Instruction::Load) {
		assert(!te->operandsvalues.empty());
		ref<Expr> addr = te->operandsvalues.front();
		if (isa<ConstantExpr> (addr)) {
			return true;
		} else
			return false;
	}
	return true;
}

int Slicer::mayWrite(TraceEntry * te) {
	if (trace_live.find(te) == trace_live.end() && isConcrete(te))
		return false;
	else
		return true;
}

void Slicer::buildDSlice(TraceEntry * exploit) {
	slice.insert(exploit);
	std::queue<TraceEntry *> tobe;
	tobe.push(exploit);
	std::set<TraceEntry *>::iterator it;
	while (!tobe.empty()) {
		std::set<TraceEntry *>::iterator sit = tobe.front()->defuse.begin();
		std::set<TraceEntry *>::iterator eit = tobe.front()->defuse.end();
		for (it = sit; it != eit; it++) {
			if ((*it) != NULL) {
				tobe.push(*it);
				slice.insert(*it);
			}
		}
		tobe.pop();
	}
}

void Slicer::printSlice(raw_ostream & out) {
	out << "-----------Output Slice size = " << slice.size()
			<< "----------------\n";
	std::set<TraceEntry *, tecomp>::iterator it;
	for (it = slice.begin(); it != slice.end(); it++) {
		printTraceEntry(*it, out);
	}
	out << "----------- End Output Slice size = " << slice.size()
			<< "----------------\n";
}

std::vector<TraceEntry *>::reverse_iterator Slicer::findCall(TraceEntry * ret,
		std::vector<TraceEntry *>::reverse_iterator start) {
	assert (ret->kins->inst->getOpcode() == Instruction::Ret);
	std::vector<TraceEntry *>::reverse_iterator rit;
	for (rit = start; rit != trace.rend(); rit++) {
		TraceEntry * te = *rit;
		if (te->kins->inst->getOpcode() == Instruction::Call
				|| te->kins->inst->getOpcode() == Instruction::Invoke) {
			return rit;
		}
	}
	return start;
}

bool SlicerExploitParser::parse(cl::Option &O, llvm::StringRef ArgName,
		llvm::StringRef Arg, unsigned long long &Val) {
	const char *ArgStart = Arg.data();
	char *End;

	// Parse integer part, leaving 'End' pointing to the first non-integer char
	Val = strtoul(ArgStart, &End, 0);

	return false;
}

void Slicer::run(TraceEntry * exploit) {
	takeToSlice(exploit); //first push the init instruction

	std::vector<TraceEntry *>::reverse_iterator rit;

	for (rit = trace.rbegin(); rit != trace.rend(); rit++) {
		TraceEntry * cur = *rit;
		//printTraceEntry(cur);
		if (cur->kins->inst->getOpcode() == Instruction::Ret) //return ins
		{
			std::vector<TraceEntry *>::reverse_iterator call = findCall(cur,
					rit);
			if (mayWriteF(*call)) {
				takeToSlice(cur);
			} else {
				rit = call;
			}
		} else if (cur->kins->inst->getOpcode() == Instruction::Invoke
				|| cur->kins->inst->getOpcode() == Instruction::Call) //callsite ins
		{
			takeToSlice(cur);
			std::set<TraceEntry *>::iterator it;
			std::set<TraceEntry *>::iterator sit = cur->defuse.begin();
			std::set<TraceEntry *>::iterator eit = cur->defuse.end();
			for (it = sit; it != eit; it++) {
				takeToSlice(cur);
			}
		} else if (cur->kins->inst->getOpcode() == Instruction::Br || //branch ins
				cur->kins->inst->getOpcode() == Instruction::Switch) //switch ins
		{
			TraceEntry * head = *(slice.begin());
			//printTraceEntry(head);
			if (!isPostDominate(head, cur) || writtenBetween(cur, head)) {
				takeToSlice(cur);
			}
		} else {
			if (mayWrite(cur)) {
				takeToSlice(cur);
			}
		}
		//printSlice();
	}
}

int Slicer::writtenBetween(TraceEntry *, TraceEntry *) {
	return true;
}
int Slicer::isPostDominate(TraceEntry * t1, TraceEntry * t2) {
	Function * func1 = t1->kins->inst->getParent()->getParent();
	Function * func2 = t2->kins->inst->getParent()->getParent();
	if (func1 == func2) {
		/*PassManager fpm;
		 AnalysisUsage *AU = new AnalysisUsage();
		 pdt = new PostDominatorTree();
		 pdt->getAnalysisUsage(*AU);
		 fpm.add(pdt);*/
		fpm->run(*func1);
		bool res = pdt->dominates(t1->kins->inst->getParent(),
				t2->kins->inst->getParent());
		//PostDominatorTree &PDT = pdt->getAnalysis<PostDominatorTree>(*func1);
		llvm::errs() << "@@@@@@@@@@@@@@@@@@@@@@" << res << "\n";
		return res;
		//return pdt->dominates(t1->kins->inst->getParent(),t2->kins->inst->getParent());
	}
	return false;
}

void Slicer::initRequiredPass(Module *m) {
	/*PassManager pm;

	 pm.add(pdt);
	 pm.run(*m);
	 pdt->print(llvm::errs(), m);*/

	AnalysisUsage *AU = new AnalysisUsage();
	pdt = new PostDominatorTree();
	pdt->getAnalysisUsage(*AU);

	fpm = new FunctionPassManager(m);
	fpm->add(pdt);
}

void Slicer::updateWritten(TraceEntry * te) {
	trace_live.erase(te);
}

void Slicer::addReadtoLive(TraceEntry * te) {
	std::set<TraceEntry *>::iterator it;
	std::set<TraceEntry *>::iterator sit = te->defuse.begin();
	std::set<TraceEntry *>::iterator eit = te->defuse.end();
	for (it = sit; it != eit; it++) {
		TraceEntry * def = *it;
		if (def != NULL)
			trace_live[def] = def->usedef;
	}
}

void Slicer::takeToSlice(TraceEntry * cur) {
	slice.insert(cur);
	updateWritten(cur);
	addReadtoLive(cur);

}

void Slicer::printTraceEntry(TraceEntry * tentry, raw_ostream & out) {
	KInstruction * kins = tentry->kins;
	Instruction *i = kins->inst;
	out << "---------------------------\n";
	out << "trace address: " << tentry << " teid: " << tentry->id << "\n";
	out << "pid:" << tentry->pid << "\t";
	out << "inst address: " << i << ": " << *i << " instance: " << tentry->instance
			<< "\n";
	std::vector<ref<Expr> >::iterator eit;
	out << "@operands are:\n";
	for (eit = tentry->operandsvalues.begin(); eit
			!= tentry->operandsvalues.end(); eit++) {
		ostringstream oss (ostringstream::out);
		(*eit)->print(oss);
		out << oss.str();
		out << " end\n";
	}
	if (!tentry->destvalue.isNull()) {
		out << "@dest is:\n";
		ostringstream oss (ostringstream::out);
		tentry->destvalue->print(oss);
		out << oss.str();
		out << " end \n";
	}
	std::set<TraceEntry *>::iterator defit;
	out << "!use!\n";
	for (defit = tentry->defuse.begin(); defit != tentry->defuse.end(); defit++) {
		TraceEntry * defentry = *defit;
		if (defentry != NULL)
			out << defentry << " teid: " << defentry->id << " ins: "
					<< defentry->kins->inst << "\n";
		else
			out << "NULL \n";
	}
	out << "!def!\n";
	for (defit = tentry->usedef.begin(); defit != tentry->usedef.end(); defit++) {
		TraceEntry * defentry = *defit;
		if (defentry != NULL)
			out << defentry << " teid: " << defentry->id << " ins: "
					<< defentry->kins->inst << "\n";
		else
			out << "NULL \n";
	}
}

TraceEntry * Slicer::readShdwmem(uint64_t addr) {
	map<uint64_t, TraceEntry *>::iterator it;
	it = shdwmem.find(addr);
	if (it == shdwmem.end()) {
		return NULL;
	} else {
		return it->second;
	}
}

void Slicer::writeShdwmem(uint64_t addr, TraceEntry * te) {
	shdwmem[addr] = te;
}

void Slicer::printTrace(raw_ostream & out) {
	std::vector<TraceEntry *>::iterator it;
	for (it = trace.begin(); it != trace.end(); it++) {
		TraceEntry * tentry = *it;
		printTraceEntry(tentry, out);
	}
}
